# -*- coding: utf-8 -*-
# @Time    : 2021/7/7 14:46
# @Author  : LuoTianHang


# ####################train.py 说明##########################
# this script is the core py to train the model
import copy

import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import torch.onnx
import torch.optim as optim
from torch import nn
import  numpy as np
from config.config import parser_config
from datalist import image_loader
from model import model
from utils import Normalization

content_layers_default = ['conv_4']
style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).cuda()
cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).cuda()


class ContentLoss(nn.Module):
    def __init__(self, target):
        super(ContentLoss, self).__init__()
        self.target = target.detach()

    def forward(self, input):
        self.loss = F.mse_loss(input, self.target)
        return input


def gram_matrix(input):
    a, b, c, d = input.size()

    features = input.view(a * b, c * d)
    G = torch.mm(features, features.t())
    return G.div(a * b * c * d)


def trans(tensor):
    image = tensor.cpu().clone()
    image = image.squeeze(0)
    image = image.numpy()
    image = image.transpose((1, 2, 0))
    return image


class StyleLoss(nn.Module):
    def __init__(self, target):
        super(StyleLoss, self).__init__()
        self.target = gram_matrix(target).detach()

    def forward(self, input):
        G = gram_matrix(input)
        self.loss = F.mse_loss(G, self.target)
        return input


class Train(object):
    def __init__(self):
        self.args = parser_config()

        print(f"-----------{self.args.project_name}-------------")

        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        if use_cuda:
            torch.cuda.manual_seed(self.args.seed)
            torch.cuda.manual_seed_all(self.args.seed)
        else:
            torch.manual_seed(self.args.seed)

        self.device = torch.device('cuda' if use_cuda else 'cpu')

        # data

        self.style_img = image_loader("data/star.jpg")
        self.content_img = image_loader("data/dog.jpg")
        self.input_img = image_loader("data/dog.jpg").cuda()

        self.input_img.requires_grad_()

        # configure the model

        normalization = Normalization(cnn_normalization_mean, cnn_normalization_std).cuda()
        self.model = nn.Sequential(normalization)
        self.content_losses = []
        self.style_losses = []

        i = 0
        for layer in model.children():

            if isinstance(layer, nn.Conv2d):
                i += 1
                name = 'conv_{}'.format(i)
            elif isinstance(layer, nn.ReLU):
                name = 'relu_{}'.format(i)
                layer=nn.ReLU(inplace=False)
            elif isinstance(layer, nn.MaxPool2d):
                name = 'pool_{}'.format(i)
            elif isinstance(layer, nn.BatchNorm2d):
                name = 'bn_{}'.format(i)

            self.model.add_module(name, layer)

            if name in content_layers_default:
                target = self.model(self.content_img).detach()
                content_loss = ContentLoss(target)
                self.model.add_module("content_loss_{}".format(i), content_loss)
                self.content_losses.append(content_loss)
            if name in style_layers_default:
                target_feature = self.model(self.style_img).detach()
                style_loss = StyleLoss(target_feature)
                self.model.add_module("style_loss_{}".format(i), style_loss)
                self.style_losses.append(style_loss)
        for i in range(len(self.model) - 1, -1, -1):
            if isinstance(self.model[i], ContentLoss) or isinstance(self.model[i], StyleLoss):
                break
        self.model = self.model[:(i + 1)]

        self.optimizer = optim.LBFGS([self.input_img])
        for i in range(40):
            self.epoch=i
            for epoch in range(self.epoch):
                self.optimizer.step(self.train)
            self.test()

    def train(self):
        self.optimizer.zero_grad()
        self.input_img.data.clamp_(0, 1)
        temp=self.input_img
        style_score = 0
        content_score = 0

        for module in self.model:
            temp = module(temp)
            if isinstance(module, ContentLoss):
                content_score = module.loss
            if isinstance(module, StyleLoss):
                style_score += module.loss

        # config the weight
        style_score *= 100000
        content_score *= 1

        loss = style_score + content_score

        loss.backward()

        return style_score + content_score


    def test(self):
        temp=self.input_img

        style_img = trans(self.style_img)
        content_img = trans(self.content_img)
        output = trans(temp.detach())

        f, ax = plt.subplots(1, 3)
        ax[0].imshow(style_img)
        ax[0].set_title('Style Image')
        ax[1].imshow(content_img)
        ax[1].set_title('Content Image')
        ax[2].imshow(output)
        ax[2].set_title('Transposed Image')
        plt.savefig("result/"+str(self.epoch)+'.png')


ii = Train()
